import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import time


# 修改后的归并排序算法
def merge_sort(arr):
    def merge(left, right, steps):
        merged = []
        i = j = 0

        while i < len(left) and j < len(right):
            if left[i] < right[j]:
                merged.append(left[i])
                i += 1
                steps.append((merged.copy(), i - 1, -1))  # 记录排序过程
            else:
                merged.append(right[j])
                j += 1
                steps.append((merged.copy(), -1, j - 1))  # 记录排序过程

        # Append any remaining elements
        while i < len(left):
            merged.append(left[i])
            i += 1
            steps.append((merged.copy(), i - 1, -1))  # 记录排序过程

        while j < len(right):
            merged.append(right[j])
            j += 1
            steps.append((merged.copy(), -1, j - 1))  # 记录排序过程

        return merged, steps

    def merge_sort_helper(arr, steps):
        if len(arr) <= 1:
            return arr, steps

        mid = len(arr) // 2
        left, steps = merge_sort_helper(arr[:mid], steps)
        right, steps = merge_sort_helper(arr[mid:], steps)
        merged, steps = merge(left, right, steps)

        return merged, steps

    steps = []
    sorted_arr, steps = merge_sort_helper(arr, steps)
    return steps


# 随机生成一个数组
arr = np.random.randint(0, 100, 20)

# 获取排序步骤
steps = merge_sort(arr.copy())

# 创建动画
fig, ax = plt.subplots()
ax.set_xlim(-1, len(arr))
ax.set_ylim(0, max(arr) + 10)

bar_rects = ax.bar(range(len(arr)), arr, align="edge")


def animate(i):
    current_step = steps[i][0]
    highlight_left = steps[i][1]
    highlight_right = steps[i][2]

    for idx, (rect, val) in enumerate(zip(bar_rects, current_step)):
        if idx == highlight_left:
            rect.set_color('r')  # 将左侧合并的柱子设为红色
        elif idx == highlight_right:
            rect.set_color('g')  # 将右侧合并的柱子设为绿色
        else:
            rect.set_color('b')  # 其余柱子设为蓝色
        rect.set_height(val)
    time.sleep(0.05)  # 可以适当调整延迟以观察排序过程

    return bar_rects


ani = animation.FuncAnimation(fig, animate, frames=len(steps), interval=50, blit=True)

plt.show()
